{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "generate a Gauss distribution for clustering\n",
    "\n",
    "mapping output of qNetwork to some distribution (measurement of probability)\n",
    "\n",
    "Oh I need to read GAN, even try some diffusion tricks.\n",
    "Target: make a Q-network to output a specific distribution? Why not Hamiltionian? \n",
    "\n",
    "1. make a bunch of data\n",
    "2. build a Q-network\n",
    "3. train the Q-network to output the distribution of the data using GAN\n",
    "4. test the Q-network\n",
    "\n",
    "WGAN: https://zhuanlan.zhihu.com/p/25071913\n",
    "1. 判别器最后一层去掉sigmoid\n",
    "2. 生成器和判别器的loss不取log\n",
    "3. 每次更新判别器的参数之后把它们的绝对值截断到不超过一个固定常数c\n",
    "4. 不要用基于动量的优化算法（包括momentum和Adam），推荐RMSProp，SGD也行"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "# import matplotlib.pyplot as plt\n",
    "\n",
    "# generation of data\n",
    "\n",
    "NumData = 1000\n",
    "\n",
    "batch_size = 1\n",
    "\n",
    "# np.random.seed(520309)\n",
    "\n",
    "# make a Dataloader\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "class GaussDataset(Dataset):\n",
    "  def __init__(self):\n",
    "    self.hists = []\n",
    "    for _ in range(NumData):\n",
    "      data = np.random.normal(loc=0, scale=1, size=1000)\n",
    "      # Create histogram\n",
    "      hist, _ = np.histogram(data, bins=8)\n",
    "      # normalization\n",
    "      hist = hist / hist.sum()\n",
    "      self.hists.append(hist)\n",
    "  def __len__(self):\n",
    "    return len(self.hists)\n",
    "  def __getitem__(self, idx):\n",
    "    return self.hists[idx]\n",
    "  \n",
    "dataset = GaussDataset()\n",
    "histdataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "# Plot histogram\n",
    "# for i_batch, sample_batched in enumerate(histdataloader):\n",
    "#   print(i_batch, sample_batched.size())\n",
    "#   for i in range(5):\n",
    "#     plt.plot(sample_batched[-i], 'o-')\n",
    "#   plt.xlabel('Bins')\n",
    "#   plt.ylabel('Frequency')\n",
    "#   plt.title('Gaussian Distribution')\n",
    "#   plt.show()\n",
    "\n",
    "print(dataset[-1])\n",
    "print('Data generation done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qiskit import QuantumRegister, QuantumCircuit\n",
    "from qiskit.circuit import QuantumCircuit, ParameterVector, Parameter\n",
    "\n",
    "qr = QuantumRegister(3, 'q')\n",
    "qc = QuantumCircuit(qr)\n",
    "\n",
    "# inputs = ParameterVector('inputs', 3)\n",
    "angles = ParameterVector('angles', 9)\n",
    "\n",
    "# qc.initialize(inputs[0], 0)\n",
    "\n",
    "# qc.ry(inputs[0], qr[0])\n",
    "# qc.ry(inputs[1], qr[1])\n",
    "# qc.ry(inputs[2], qr[2])\n",
    "\n",
    "# Layer 0\n",
    "qc.h(qr[0:3])\n",
    "qc.ry(angles[0], qr[0])\n",
    "qc.ry(angles[1], qr[1])\n",
    "qc.ry(angles[2], qr[2])\n",
    "\n",
    "# Layer 1\n",
    "\n",
    "qc.cx(qr[0], qr[1])\n",
    "qc.ry(angles[3], qr[0])\n",
    "qc.cx(qr[1], qr[2])\n",
    "qc.ry(angles[4], qr[1])\n",
    "qc.ry(angles[5], qr[2])\n",
    "\n",
    "# Layer 2\n",
    "qc.cx(qr[0], qr[1])\n",
    "qc.ry(angles[6], qr[0])\n",
    "qc.cx(qr[1], qr[2])\n",
    "qc.ry(angles[7], qr[1])\n",
    "qc.ry(angles[8], qr[2])\n",
    "\n",
    "# bind the parameters\n",
    "# qc = qc.bind_parameters({\n",
    "#   angles[0]: 0.5, angles[1]: 0.5, angles[2]: 0.5, angles[3]: 0.5, angles[4]: 0.5, angles[5]: 0.5, angles[6]: 0.5, angles[7]: 0.5, angles[8]: 0.5})\n",
    "\n",
    "# qc.measure_all()\n",
    "\n",
    "# qc.draw('mpl')\n",
    "print('Circuit done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build a GAN discriminator\n",
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, input_size):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.leaky_relu = nn.LeakyReLU(0.2)\n",
    "        # self.linear_input = nn.Linear(input_size, 5)\n",
    "        self.linear20 = nn.Linear(input_size, 1)\n",
    "        # WGAN: remove sigmoid\n",
    "        # self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, input: torch.Tensor) -> torch.Tensor:\n",
    "        # x = self.linear_input(input)\n",
    "        x = self.leaky_relu(input)\n",
    "        x = self.linear20(x)\n",
    "        # x = self.sigmoid(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# pytorch QNN\n",
    "from qiskit_machine_learning.neural_networks import SamplerQNN\n",
    "from qiskit_algorithms.utils import algorithm_globals\n",
    "from qiskit_machine_learning.connectors import TorchConnector\n",
    "\n",
    "from qiskit_aer.primitives import Sampler\n",
    "\n",
    "# from qiskit.primitives import Sampler\n",
    "\n",
    "\n",
    "shots = 8192\n",
    "\n",
    "run_options = {\n",
    "    'shots': shots\n",
    "}\n",
    "\n",
    "# if cuda is available, use it\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "backend_options = {\n",
    "    'method': 'statevector',\n",
    "    'device': 'GPU' if device.type == 'cuda' else 'CPU',\n",
    "    'max_job_size': 100, # maximum number of circuits in one job\n",
    "    'max_parallel_threads': 4, # maximum number of CPU cores used by OpenMP for parallelization\n",
    "    'max_parallel_experiments': 8192, # maximum number of experiments for parallelization\n",
    "}\n",
    "\n",
    "sampler = Sampler(run_options=run_options, backend_options=backend_options)\n",
    "\n",
    "print('Sampler done using', backend_options.get('device'))\n",
    "\n",
    "def create_generator() -> TorchConnector:\n",
    "    qnn = SamplerQNN(\n",
    "        circuit=qc,\n",
    "        sampler=sampler,\n",
    "        input_params=[],\n",
    "        weight_params=qc.parameters,\n",
    "        sparse=False,\n",
    "    )\n",
    "\n",
    "    initial_weights = algorithm_globals.random.random(qc.num_parameters)\n",
    "    return TorchConnector(qnn, initial_weights)\n",
    "\n",
    "\n",
    "generator = create_generator()\n",
    "discriminator = Discriminator(8)\n",
    "\n",
    "criterion = nn.BCELoss()\n",
    "\n",
    "\n",
    "from torch.optim import Adam, RMSprop\n",
    "\n",
    "lr = 0.0002  # learning rate\n",
    "b1 = 0.5  # first momentum parameter\n",
    "b2 = 0.999  # second momentum parameter\n",
    "\n",
    "#WGAN: change to RMSprop\n",
    "generator_optimizer = RMSprop(generator.parameters(), lr=lr, weight_decay=0.005)\n",
    "discriminator_optimizer = RMSprop(\n",
    "    discriminator.parameters(), lr=lr, weight_decay=0.005\n",
    ")\n",
    "# generator_optimizer = Adam(generator.parameters(), lr=lr, betas=(b1, b2), weight_decay=0.005)\n",
    "# discriminator_optimizer = Adam(\n",
    "#     discriminator.parameters(), lr=lr, betas=(b1, b2), weight_decay=0.005\n",
    "# )\n",
    "\n",
    "\n",
    "\n",
    "# gen_dist = generator(torch.tensor([[] for _ in range(batch_size)]))\n",
    "# print(gen_dist)\n",
    "\n",
    "print('GAN done')\n",
    "print('using device:', device)\n",
    "\n",
    "# from torchsummary import summary\n",
    "# summary(discriminator,input_size=(8,))\n",
    "# # summary(generator,input_size=(0,))\n",
    "# print(generator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "starttime = time.strftime(\"%Y-%m-%d_%H_%M_%S\")\n",
    "print(\"Start experiment:\", starttime)#定义实验时间\n",
    "random_seed = str(np.random.get_state()[1][0])\n",
    "print('Random seed:', random_seed)\n",
    "writer = SummaryWriter(log_dir=\"./runs/\"+starttime+random_seed,comment=starttime+random_seed,flush_secs=60)\n",
    "\n",
    "generator.to(device)\n",
    "discriminator.to(device)\n",
    "\n",
    "def relative_entropy(p, q):\n",
    "    # avoid nan and inf\n",
    "    p = torch.where(p <= 1e-6, torch.tensor(1e-6), p)\n",
    "    q = torch.where(q <= 1e-6, torch.tensor(1e-6), q)\n",
    "    return (p * torch.log(p / q)).sum()\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "n_epochs = 300\n",
    "min_entropy = 1e6\n",
    "\n",
    "start = time.time()\n",
    "print('Start training')\n",
    "for epoch in range(n_epochs):\n",
    "    generator_loss_values = []\n",
    "    discriminator_loss_values = []\n",
    "    entropy_values = []\n",
    "    print(f\"Epoch {epoch}\", time.time() - start)\n",
    "    for i_batch, sample_batched in tqdm(enumerate(histdataloader), total=len(histdataloader)):\n",
    "        # print(i_batch, sample_batched.size())\n",
    "        # update D network: maximize log(D(x)) + log(1 - D(G(z)))\n",
    "        # print(sample_batched.shape[0])\n",
    "        \n",
    "        \n",
    "        ## train with real\n",
    "        # sample_batched.to(device)\n",
    "        # print(\"train with real\", time.time() - start)\n",
    "        discriminator_optimizer.zero_grad()\n",
    "        real_dist = sample_batched.to(device=device, dtype=torch.float)\n",
    "        # real_dist = torch.tensor(sample_batched, dtype=torch.float, device=device)\n",
    "\n",
    "        # real_dist = torch.where(real_dist <= 1e-6, torch.tensor(1e-6), real_dist)\n",
    "        \n",
    "        discriminat_real_dist = discriminator(real_dist)\n",
    "        # print(\"real\", real_dist.shape, discriminat_real_dist.shape)\n",
    "        #  WGAN: remove log in criterion\n",
    "        real_loss = torch.mean(discriminat_real_dist)\n",
    "        # real_loss = criterion(discriminat_real_dist, torch.ones(sample_batched.shape[0], 1, device=device))\n",
    "        real_loss.backward()\n",
    "        # print(\"train with real\", time.time() - start)\n",
    "\n",
    "        ## train with fake\n",
    "        # print(\"train with fake\", time.time() - start)\n",
    "        gen_dist = generator(torch.tensor([[] for _ in range(sample_batched.shape[0])], device=device))\n",
    "        # gen_dist = torch.where(gen_dist <= 1e-6, torch.tensor(1e-6), gen_dist)\n",
    "        discriminat_gen_dist = discriminator(gen_dist.detach())\n",
    "        # print(\"gen\", gen_dist.shape, discriminat_gen_dist.shape)\n",
    "        #  WGAN: remove log in criterion\n",
    "        fake_loss = -torch.mean(discriminat_gen_dist)\n",
    "        # fake_loss = criterion(discriminat_gen_dist, torch.zeros(sample_batched.shape[0], 1, device=device))\n",
    "        fake_loss.backward()\n",
    "        discriminator_optimizer.step()\n",
    "        # print(\"train with fake\", time.time() - start)\n",
    "\n",
    "        # Clip weights of discriminator\n",
    "        for p in discriminator.parameters():\n",
    "            p.data.clamp_(-100, 100)\n",
    "\n",
    "        # update G network: maximize log(D(G(z)))\n",
    "        # print(\"update G network\", time.time() - start)\n",
    "        generator_optimizer.zero_grad()\n",
    "        # print(discriminator(gen_dist))\n",
    "        #  WGAN: remove log in criterion\n",
    "        generator_loss = torch.mean(discriminator(gen_dist))\n",
    "        # generator_loss = criterion(discriminator(gen_dist), torch.ones(sample_batched.shape[0], 1, device=device))\n",
    "        generator_loss.backward()\n",
    "        generator_optimizer.step()\n",
    "        # print(\"update G network\", time.time() - start)\n",
    "\n",
    "        # # relative entropy\n",
    "        mean_entropy = relative_entropy(real_dist, gen_dist.detach()).mean()\n",
    "        # # print(f\"Epoch {epoch}, Batch {i_batch}, D Loss: {real_loss + fake_loss}, G Loss: {generator_loss}, Entropy: {mean_entropy}, real: {real_dist}, gen: {gen_dist}\")\n",
    "\n",
    "        # # if np.isnan(mean_entropy) or np.isinf(mean_entropy):\n",
    "        # #     print(f\"Epoch {epoch}, Batch {i_batch}, D Loss: {real_loss + fake_loss}, G Loss: {generator_loss}, Entropy: {entropy(real_dist, gen_dist.detach()).mean()}, real: {real_dist}, gen: {gen_dist}\")\n",
    "\n",
    "        batch_log_index = epoch * NumData / batch_size + i_batch\n",
    "                \n",
    "        writer.add_scalar('BatchLoss/DLoss', generator_loss.detach().item(), batch_log_index)\n",
    "        writer.add_scalar('BatchLoss/GLoss', real_loss.detach() + fake_loss.detach(), batch_log_index)\n",
    "        writer.add_scalar('Batch/Entropy', mean_entropy, batch_log_index)\n",
    "\n",
    "        generator_loss_values.append(generator_loss.detach().item())\n",
    "        discriminator_loss_values.append(real_loss.detach().item() + fake_loss.detach().item())\n",
    "        entropy_values.append(mean_entropy.detach().item())\n",
    "        \n",
    "    writer.add_scalar('Loss/DLoss', np.mean(discriminator_loss_values), epoch)\n",
    "    writer.add_scalar('Loss/GLoss', np.mean(generator_loss_values), epoch)\n",
    "    writer.add_scalar('Entropy', np.mean(entropy_values), epoch)\n",
    "\n",
    "    # save the model with least entropy\n",
    "    if epoch == 0 or np.mean(entropy_values) < min_entropy:\n",
    "        min_entropy = np.mean(entropy_values)\n",
    "        torch.save(generator.state_dict(), f'generator-{random_seed}.pth')\n",
    "        torch.save(discriminator.state_dict(), f'discriminator-{random_seed}.pth')\n",
    "        print(f\"Save model on Epoch {epoch}, D Loss: {np.mean(discriminator_loss_values)}, G Loss: {np.mean(generator_loss_values)}, Entropy: {np.mean(entropy_values)}, real: {real_dist}, gen: {gen_dist}\")\n",
    "\n",
    "\n",
    "\n",
    "elapsed = time.time() - start\n",
    "print(f\"Fit in {elapsed:0.2f} sec\")\n",
    "\n",
    "# plot the result, hist graph for real_dist and line dot for gen_dist\n",
    "# real_dist, gen_dist\n",
    "# plt.bar(bins[:-1], real_dist.detach().numpy(), width=(bins[1]-bins[0]))\n",
    "# plt.plot(bins[:-1], gen_dist.detach().numpy(), 'r')\n",
    "# plt.show()\n",
    "# plt.bar(bins[:-1], real_dist.detach().numpy(), width=(bins[1]-bins[0]))\n",
    "# plt.bar(bins[:-1], gen_dist.detach().numpy(), width=(bins[1]-bins[0]))\n",
    "# plt.show()\n",
    "# print(hist)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the model\n",
    "generator = create_generator()\n",
    "generator.load_state_dict(torch.load(f'generator-{random_seed}.pth', map_location=torch.device('cpu')))\n",
    "generator.eval()\n",
    "\n",
    "# generate the distribution\n",
    "gen_dist = generator(torch.tensor([[] for _ in range(125)])).detach().numpy()\n",
    "# print(gen_dist)\n",
    "# convert gen_dist 1D\n",
    "# gen_dist = gen_dist * 1000\n",
    "# data = np.random.normal(loc=3.5, scale=1, size=1000)\n",
    "# normalize\n",
    "print('Model loaded', gen_dist.shape)\n",
    "\n",
    "# merge gen_dist into one average\n",
    "gen_dist = gen_dist.mean(axis=0)\n",
    "dataset_changed = np.array(dataset).mean(axis=0)\n",
    "\n",
    "# Plot 10 histogram\n",
    "import matplotlib.pyplot as plt\n",
    "plt.plot(gen_dist, 'o-g', label='gen')\n",
    "# plt.hist(data, bins=8, alpha=0.5, label='real',range=(0, 7))\n",
    "plt.bar(range(len(dataset_changed)), dataset_changed, label='real', alpha=0.5)\n",
    "plt.xlabel('Bins')\n",
    "plt.ylabel('Frequency')\n",
    "plt.title('Gaussian Distribution')\n",
    "plt.legend()\n",
    "# plt.show()\n",
    "plt.savefig(f'hist-{random_seed}.png')\n",
    "# print('Plot done')\n",
    "\n",
    "# print(gen_dist)\n",
    "# draw normal distribution using plt.hist\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Why we design the circuit in this way?\n",
    "\n",
    "# superposition + parameter ?\n",
    "\n",
    "# How to back propagate the parameter?\n",
    "\n",
    "# How to use the quantum feature to detect the minimum of the function?\n",
    "\n",
    "# Gedeer Encoding\n",
    "\n",
    "# EfficientSU2 ?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qiskit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
